/*
 * Copyright 2022 Google LLC
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#include "src/gpu/graphite/dawn/DawnBuffer.h"

#include "include/private/base/SkAlign.h"
#include "src/gpu/graphite/Log.h"
#include "src/gpu/graphite/dawn/DawnAsyncWait.h"
#include "src/gpu/graphite/dawn/DawnSharedContext.h"

namespace skgpu::graphite {
#ifdef SK_DEBUG
static const char* kBufferTypeNames[kBufferTypeCount] = {
        "Vertex",
        "Index",
        "Xfer CPU to GPU",
        "Xfer GPU to CPU",
        "Uniform",
        "Storage",
        "Indirect",
        "VertexStorage",
        "IndexStorage",
};
#endif

sk_sp<DawnBuffer> DawnBuffer::Make(const DawnSharedContext *sharedContext, size_t size, BufferType type,
    AccessPattern accessPattern)
{
    return DawnBuffer::Make(sharedContext, size, type, accessPattern,
#ifdef SK_DEBUG
        /* label= */ kBufferTypeNames[static_cast<int>(type)]
#else
        /* label= */ nullptr
#endif
    );
}

sk_sp<DawnBuffer> DawnBuffer::Make(const DawnSharedContext *sharedContext, size_t size, BufferType type,
    AccessPattern accessPattern, const char *label)
{
    if (size <= 0) {
        return nullptr;
    }

    wgpu::BufferUsage usage = wgpu::BufferUsage::None;

    switch (type) {
        case BufferType::kVertex:
            usage = wgpu::BufferUsage::Vertex | wgpu::BufferUsage::CopyDst;
            break;
        case BufferType::kIndex:
            usage = wgpu::BufferUsage::Index | wgpu::BufferUsage::CopyDst;
            break;
        case BufferType::kXferCpuToGpu:
            usage = wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::MapWrite;
            break;
        case BufferType::kXferGpuToCpu:
            usage = wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapRead;
            break;
        case BufferType::kUniform:
            usage = wgpu::BufferUsage::Uniform | wgpu::BufferUsage::CopyDst;
            break;
        case BufferType::kStorage:
            usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::CopySrc;
            break;
        case BufferType::kIndirect:
            usage = wgpu::BufferUsage::Indirect | wgpu::BufferUsage::Storage;
            break;
        case BufferType::kVertexStorage:
            usage = wgpu::BufferUsage::Vertex | wgpu::BufferUsage::Storage;
            break;
        case BufferType::kIndexStorage:
            usage = wgpu::BufferUsage::Index | wgpu::BufferUsage::Storage;
            break;
    }

    if (sharedContext->caps()->drawBufferCanBeMapped() && accessPattern == AccessPattern::kHostVisible &&
        type != BufferType::kXferGpuToCpu) {
        usage |= wgpu::BufferUsage::MapWrite;
    }

    wgpu::BufferDescriptor desc;
    desc.label = label;
    desc.usage = usage;
    desc.size = size;
    // Specifying mappedAtCreation avoids clearing the buffer on the GPU which can cause MapAsync to
    // be very slow as it waits for GPU execution to complete.
    desc.mappedAtCreation = SkToBool(usage & wgpu::BufferUsage::MapWrite);

    auto buffer = sharedContext->device().CreateBuffer(&desc);
    if (!buffer) {
        return {};
    }

    void *mappedAtCreationPtr = nullptr;
    if (desc.mappedAtCreation) {
        mappedAtCreationPtr = buffer.GetMappedRange();
        SkASSERT(mappedAtCreationPtr);
    }

    return sk_sp<DawnBuffer>(new DawnBuffer(sharedContext, size, std::move(buffer), mappedAtCreationPtr));
}

void DawnBuffer::prepareForReturnToCache(const std::function<void()> &takeRef)
{
    // This implementation is almost Dawn-agnostic. However, Buffer base class doesn't have any
    // way of distinguishing a buffer that is mappable for writing from one mappable for reading.
    // We only need to re-map the former.
    if (!(fBuffer.GetUsage() & wgpu::BufferUsage::MapWrite)) {
        return;
    }
    // We cannot start an async map while the GPU is still using the buffer. We asked that
    // our Resource convert command buffer refs to usage refs. So we should never have any
    // command buffer refs.
    SkASSERT(!this->debugHasCommandBufferRef());
    // Note that the map state cannot change on another thread when we are here. We got here
    // because there were no UsageRefs on the buffer but async mapping holds a UsageRef until it
    // completes.
    if (this->isMapped()) {
        return;
    }
    takeRef();
    this->asyncMap(
        [](void *ctx, skgpu::CallbackResult result) {
            sk_sp<DawnBuffer> buffer(static_cast<DawnBuffer *>(ctx));
            if (result != skgpu::CallbackResult::kSuccess) {
                buffer->setDeleteASAP();
            }
        },
        this);
}

DawnBuffer::DawnBuffer(const DawnSharedContext *sharedContext, size_t size, wgpu::Buffer buffer,
    void *mappedAtCreationPtr)
    : Buffer(sharedContext, size,
    /* commandBufferRefsAsUsageRefs= */ buffer.GetUsage() & wgpu::BufferUsage::MapWrite),
      fBuffer(std::move(buffer))
{
    fMapPtr = mappedAtCreationPtr;
}

void DawnBuffer::onAsyncMap(GpuFinishedProc proc, GpuFinishedContext ctx)
{
    if (proc) {
        SkAutoMutexExclusive ex(fAsyncMutex);
        if (this->isMapped()) {
            proc(ctx, CallbackResult::kSuccess);
            return;
        }
        fAsyncMapCallbacks.push_back(RefCntedCallback::Make(proc, ctx));
    }
    if (this->isUnmappable()) {
        return;
    }
    SkASSERT(fBuffer);
    SkASSERT((fBuffer.GetUsage() & wgpu::BufferUsage::MapRead) || (fBuffer.GetUsage() & wgpu::BufferUsage::MapWrite));
    SkASSERT(fBuffer.GetMapState() == wgpu::BufferMapState::Unmapped);
    bool isWrite = fBuffer.GetUsage() & wgpu::BufferUsage::MapWrite;
    auto buffer = sk_ref_sp(this);

    fBuffer.MapAsync(
        isWrite ? wgpu::MapMode::Write : wgpu::MapMode::Read, 0, fBuffer.GetSize(),
        [](WGPUBufferMapAsyncStatus s, void *userData) {
            sk_sp<DawnBuffer> buffer(static_cast<DawnBuffer *>(userData));
            SkAutoMutexExclusive em(buffer->fAsyncMutex);
            if (s == WGPUBufferMapAsyncStatus_Success) {
                if (buffer->fBuffer.GetUsage() & wgpu::BufferUsage::MapWrite) {
                    buffer->fMapPtr = buffer->fBuffer.GetMappedRange();
                } else {
                    // If buffer is only created with MapRead usage, Dawn only allows returning
                    // constant pointer. We need to use const_cast as a workaround here.
                    buffer->fMapPtr = const_cast<void *>(buffer->fBuffer.GetConstMappedRange());
                }
            } else {
                const char *statusStr;
                Priority priority = Priority::kError;
                switch (s) {
                    case WGPUBufferMapAsyncStatus_ValidationError:
                        statusStr = "ValidationError";
                        break;
                    case WGPUBufferMapAsyncStatus_Unknown:
                        statusStr = "Unknown";
                        break;
                    case WGPUBufferMapAsyncStatus_DeviceLost:
                        statusStr = "DeviceLost";
                        break;
                    case WGPUBufferMapAsyncStatus_DestroyedBeforeCallback:
                        statusStr = "DestroyedBeforeCallback";
                        priority = Priority::kDebug;
                        break;
                    case WGPUBufferMapAsyncStatus_UnmappedBeforeCallback:
                        statusStr = "UnmappedBeforeCallback";
                        priority = Priority::kDebug;
                        break;
                    case WGPUBufferMapAsyncStatus_MappingAlreadyPending:
                        statusStr = "MappingAlreadyPending";
                        break;
                    case WGPUBufferMapAsyncStatus_OffsetOutOfRange:
                        statusStr = "OffsetOutOfRange";
                        break;
                    case WGPUBufferMapAsyncStatus_SizeOutOfRange:
                        statusStr = "SizeOutOfRange";
                        break;
                    default:
                        statusStr = "<Other>";
                        break;
                }
                SKGPU_LOG(priority, "Buffer async map failed with status %s.", statusStr);
                for (auto &cb : buffer->fAsyncMapCallbacks) {
                    cb->setFailureResult();
                }
            }
            buffer->fAsyncMapCallbacks.clear();
        },
        buffer.release());
}

void DawnBuffer::onMap()
{
    SKGPU_LOG_W("Synchronous buffer mapping not supported in Dawn. Failing map request.");
}

void DawnBuffer::onUnmap()
{
    SkASSERT(fBuffer);
    SkASSERT(this->isUnmappable());

    fMapPtr = nullptr;
    fBuffer.Unmap();
}

bool DawnBuffer::isUnmappable() const
{
    return fBuffer.GetMapState() != wgpu::BufferMapState::Unmapped;
}

void DawnBuffer::freeGpuData()
{
    fBuffer = nullptr;
}
} // namespace skgpu::graphite
